import numpy as np
import pyqtgraph as pg
import pyqtgraph.opengl as gl
from pyqtgraph.Qt import QtCore, QtWidgets

class MeshBuilder:
    """
    一个简单的 3D 几何体构建辅助类。
    用于将简单的几何形状（方块、锥体）拼接成复杂的网格。
    """
    def __init__(self):
        self.verts = []
        self.faces = []
    
    def add_box(self, center, size):
        """添加一个长方体"""
        x, y, z = center
        dx, dy, dz = size[0]/2, size[1]/2, size[2]/2
        
        # 记录当前顶点起始索引
        start_idx = len(self.verts)
        
        # 8个顶点
        v = [
            [x-dx, y-dy, z-dz], [x+dx, y-dy, z-dz], [x+dx, y+dy, z-dz], [x-dx, y+dy, z-dz], # 底面 0-3
            [x-dx, y-dy, z+dz], [x+dx, y-dy, z+dz], [x+dx, y+dy, z+dz], [x-dx, y+dy, z+dz]  # 顶面 4-7
        ]
        self.verts.extend(v)
        
        # 12个面 (每个面由2个三角形组成)
        # 顺时针或逆时针并不太影响 pyqtgraph，因为它通常双面渲染
        f_indices = [
            [0,1,2], [0,2,3], # Bottom
            [4,5,6], [4,6,7], # Top
            [0,1,5], [0,5,4], # Front (Y-)
            [2,3,7], [2,7,6], # Back (Y+)
            [1,2,6], [1,6,5], # Right (X+)
            [3,0,4], [3,4,7]  # Left (X-)
        ]
        
        # 将局部索引转换为全局索引
        self.faces.extend([[idx + start_idx for idx in tri] for tri in f_indices])

    def add_pyramid(self, center, size, axis='x'):
        """添加一个简单的锥体 (用于机头或导弹头)"""
        x, y, z = center
        dx, dy, dz = size[0]/2, size[1]/2, size[2]/2
        start_idx = len(self.verts)
        
        # 简单起见，底面是矩形，顶点汇聚于一点
        # 这里假设 axis='x'，即尖端指向 X 正方向
    
        tip = [x+dx, y, z] # 尖端
        # 底面 4 点
        base = [
            [x-dx, y-dy, z-dz], 
            [x-dx, y+dy, z-dz], 
            [x-dx, y+dy, z+dz], 
            [x-dx, y-dy, z+dz]
        ]
        
        self.verts.append(tip)
        self.verts.extend(base)
        
        # 0是尖端，1-4是底面
        f_indices = [
            [1,2,3], [1,3,4], # 底面封闭
            [0,1,2], # 侧面1
            [0,2,3], # 侧面2
            [0,3,4], # 侧面3
            [0,4,1]  # 侧面4
        ]
        self.faces.extend([[idx + start_idx for idx in tri] for tri in f_indices])

    def get_mesh_data(self):
        """返回 pyqtgraph.opengl.MeshData 对象"""
        return gl.MeshData(vertexes=np.array(self.verts), faces=np.array(self.faces))

import trimesh
import numpy as np
import pyqtgraph.opengl as gl

def load_and_simplify_model(file_path, target_faces=300, scale=1.0):
    """
    加载 STL/OBJ 模型，减少面数，并转换为 pyqtgraph 可用的 MeshData
    :param file_path: 模型路径 (.obj 或 .stl)
    :param target_faces: 目标面数 (越低性能越好，推荐 100-500)
    :param scale: 缩放比例
    :return: gl.MeshData
    """
    # 1. 使用 trimesh 加载模型
    mesh = trimesh.load(file_path)
    
    # 如果是包含多个物体的场景，合并成一个 mesh
    if isinstance(mesh, trimesh.Scene):
        mesh = trimesh.util.concatenate(mesh.dump())

    # 2. 减面 (Decimation) - 这一步对性能至关重要
    # 如果当前面数已经很少，就不减了
    if len(mesh.faces) > target_faces:
        print(f"原始面数: {len(mesh.faces)}，正在简化为 {target_faces}...")
        # 这种简化算法会保持物体的大致形状
        mesh = mesh.simplify_quadratic_decimation(target_faces)
        print(f"简化后面数: {len(mesh.faces)}")
    
    # 3. 提取顶点和面
    # trimesh 的顶点通常是 (N, 3) 的 float 数组
    verts = mesh.vertices * scale
    faces = mesh.faces

    # 4. 生成 pyqtgraph 的 MeshData
    # 计算法线可以让光照效果更好
    mesh_data = gl.MeshData(vertexes=verts, faces=faces)
    
    return mesh_data

# --- 使用示例 ---
# 假设你下载了一个 'drone.obj'
# mesh_data = load_and_simplify_model("drone.obj", target_faces=200, scale=0.1)
# item = gl.GLMeshItem(meshdata=mesh_data, shader='shaded', color=(0.8, 0.8, 0.8, 1))
# view.addItem(item)

class AircraftFactory:
    """
    飞机模型工厂
    X轴正方向为机头方向
    Z轴正方向为上方
    """
    
    @staticmethod
    def get_jet():
        """生成的战斗机模型 (三角翼)"""
        mb = MeshBuilder()
        # 机身 (长条)
        mb.add_box((0, 0, 0), (1.5, 0.25, 0.25))
        # 进气道/机身加宽
        mb.add_box((-0.2, 0, 0), (0.8, 0.5, 0.2))
        # 机翼 (向后掠的扁平块)
        # 用多个 Box 模拟三角翼形状
        mb.add_box((-0.3, 0, 0), (0.6, 1.4, 0.05)) # 主翼
        # 垂直尾翼
        mb.add_box((-0.6, 0, 0.3), (0.4, 0.05, 0.6))
        # 平尾
        mb.add_box((-0.7, 0, 0), (0.3, 0.8, 0.05))
        return mb.get_mesh_data()

    @staticmethod
    def get_propeller_plane():
        """生成的螺旋桨飞机 (类似塞斯纳)"""
        mb = MeshBuilder()
        # 机身
        mb.add_box((0, 0, 0), (1.2, 0.3, 0.3))
        # 机尾变细的部分 (简化为连接的一个较小的Box)
        mb.add_box((-0.8, 0, 0.05), (0.6, 0.15, 0.15))
        # 主翼 (上单翼)
        mb.add_box((0.2, 0, 0.2), (0.3, 1.8, 0.05))
        # 垂直尾翼
        mb.add_box((-1.0, 0, 0.2), (0.3, 0.05, 0.4))
        # 水平尾翼
        mb.add_box((-1.0, 0, 0.05), (0.25, 0.8, 0.05))
        # 螺旋桨 (十字)
        mb.add_box((0.6, 0, 0), (0.02, 0.8, 0.05))
        mb.add_box((0.6, 0, 0), (0.02, 0.05, 0.8))
        return mb.get_mesh_data()

    @staticmethod
    def get_flying_wing():
        """生成的飞翼/隐身无人机 (B2风格)"""
        mb = MeshBuilder()
        # 中心机身
        mb.add_box((0, 0, 0), (0.8, 0.4, 0.15))
        # 机翼 (大后掠角) - 这里用阶梯状的 Box 模拟 V 形
        mb.add_box((-0.2, 0.6, 0), (0.6, 0.8, 0.08)) # 左内
        mb.add_box((-0.2, -0.6, 0), (0.6, 0.8, 0.08)) # 右内
        mb.add_box((-0.5, 1.2, 0), (0.4, 0.6, 0.05)) # 左外
        mb.add_box((-0.5, -1.2, 0), (0.4, 0.6, 0.05)) # 右外
        return mb.get_mesh_data()

    @staticmethod
    def get_rocket():
        """生成的火箭/导弹"""
        mb = MeshBuilder()
        # 弹体
        mb.add_box((0, 0, 0), (2.0, 0.2, 0.2))
        # 弹头 (锥体) - 这里用pyqtgraph自带的 Cylinder 可能更圆，但为了统一这里用 Box 模拟或自定义 Pyramid
        mb.add_pyramid((1.1, 0, 0), (0.4, 0.2, 0.2), axis='x')
        # 尾翼 (X型)
        fin_size = 0.6
        mb.add_box((-0.8, 0, 0), (0.3, fin_size, 0.05)) # 水平
        mb.add_box((-0.8, 0, 0), (0.3, 0.05, fin_size)) # 垂直
        return mb.get_mesh_data()
    
    @staticmethod
    def get_speed():
        return load_and_simplify_model('resource/3DModels/craft_speederC.obj')

class ModelShowcase(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("3D Aircraft Models for Swarm Simulation")
        self.resize(1200, 600)
        
        # 布局
        w = QtWidgets.QWidget()
        self.setCentralWidget(w)
        layout = QtWidgets.QHBoxLayout()
        w.setLayout(layout)
        
        # 定义模型列表
        models = [
            ("Jet Fighter", AircraftFactory.get_jet(), (0.7, 0.7, 0.7, 1)),
            ("Propeller Plane", AircraftFactory.get_propeller_plane(), (0.2, 0.8, 0.2, 1)),
            ("Flying Wing", AircraftFactory.get_flying_wing(), (0.1, 0.1, 0.1, 1)),
            ("Missile/Rocket", AircraftFactory.get_rocket(), (1, 0.4, 0, 1)),
            ("Speed", AircraftFactory.get_speed(), (0.1, 0.4, 0.8, 0.3))
        ]
        
        self.views = []
        self.items = []
        
        for name, mesh_data, color in models:
            # 每个模型创建一个 View
            view = gl.GLViewWidget()
            view.opts['distance'] = 5
            view.setWindowTitle(name)
            view.setBackgroundColor(100, 100, 100)
            
            # 添加网格
            g = gl.GLGridItem()
            view.addItem(g)
            
            # 添加模型
            # shader='shaded' 让平面看起来有光影，这对于低多边形模型很重要
            # smooth=False 保持硬朗的低多边形风格
            item = gl.GLMeshItem(meshdata=mesh_data, smooth=False, drawEdges=True, 
                                 edgeColor=(0.4, 0.9, 1.0, 1.0), shader='shaded', color=color)
            item.setGLOptions('additive')
            view.addItem(item)
            
            # 添加标签 (文字)
            # 注意: GLTextItem 在某些环境可能性能一般，仅用于调试展示
            text = gl.GLTextItem(pos=(0, 0, 2), text=name, color=(1,1,1,1))
            view.addItem(text)
            
            layout.addWidget(view)
            self.views.append(view)
            self.items.append(item)

        # 动画定时器
        self.timer = QtCore.QTimer()
        self.timer.timeout.connect(self.animate)
        self.timer.start(20)

    def animate(self):
        # 让所有模型旋转展示
        for item in self.items:
            item.rotate(1, 0, 0, 1) # 绕 Z 轴旋转

if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)
    win = ModelShowcase()
    win.show()
    sys.exit(app.exec())